Multimodal SSL pretraining (MAE, SimMIM, I-JEPA, V-JEPA) with multi-arch backbones - #45
Draft
Rian354 wants to merge 1 commit into
Draft
Multimodal SSL pretraining (MAE, SimMIM, I-JEPA, V-JEPA) with multi-arch backbones#45Rian354 wants to merge 1 commit into
Rian354 wants to merge 1 commit into
Conversation
…-arch backbones Add pyhealth/models/pretrain/, a self-supervised pretraining stack for the multimodal clinical-sequence encoder. build_backbone selects among transformer, jamba and mamba backbones behind a uniform (x, mask) -> (emb, cls) contract, so every objective composes with every backbone. The mamba layer pads masked positions so its causal conv cannot leak across the mask, and RoPE is available for long-sequence extrapolation. Four objectives share the unified embedding and one mask generator: MAE (decoder reconstruction), SimMIM (masked-token linear head), and I-JEPA / V-JEPA (EMA-target latent prediction, with V-JEPA adding multi-scale target blocks). PretrainTrainer is DDP-native, supports AMP and gradient accumulation, drives the EMA target encoder for the JEPA variants, and skips DistributedSampler for IterableDataset streams that already shard internally. It resumes atomically from last.ckpt plus _resume.pt, and warm-resumes from weights alone when optimizer state is missing, so preempted cluster jobs continue instead of restarting. Per-modality losses are tracked to metrics_history.json and W&B. scripts/pretrain_ssl.py drives it (--arch / --method / --task, standardized 128/2/4 encoders, notes_only via include_labs=False). run_full_pretrain.py replays tuned hyperparameters from a best_params JSON and can override them with --extra. SLURM and Condor runners skip only encoders that reached a full 50 epochs, so partial runs resume, and set expandable_segments to avoid fragmentation OOM. W&B logging is env-gated, namespaced and auto-tagged by arch/method/task. Includes unit tests for masking, all four objectives, backbone contracts, RoPE, and checkpoint transfer into the downstream Transformer.
Rian354
marked this pull request as draft
August 4, 2026 06:05
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
pyhealth/models/pretrain/, a self-supervised pretraining stack for the multimodal clinical-sequence encoder, with four SSL objectives sharing one unified embedding and mask generatorbuild_backbone(arch)selectstransformer/jamba/mambabehind a single(x, mask) -> (emb, cls)contract, so every objective works with every backbonePretrainTrainer, DDP-native with atomic checkpoint resume, so preempted cluster jobs continue instead of restarting from zeroscripts/pretrain_ssl.pyplus SLURM/Condor runners for full-scale 50-epoch runs at tuned hyperparametersBackbones
transformerTransformerLayeruse_rope,rope_scalingfor NTK-aware extrapolation)mambaMambaLayer(new)MambaBlock+RMSNorm; pads masked positions to 0 so the causal conv cannot leak across the maskjambaJambaLayernum_transformer_layers/num_mamba_layersSSL objectives
TransformerLayerdecoder reconstructs masked tokens;target='token'(content-only, per-embedding normalized) or'unified'(2,4,8)with ascale_embed, and per-scale loss breakdownUnifiedMaskGeneratorsupports random and block masking withper_modality_ratiooverrides; block masking preserves contiguity and enforces a minimum block length.PretrainTrainer
RANK/WORLD_SIZEfrom torchrun; AMP and gradient accumulation; EMA hook everyema_update_everyoptimizer steps for the JEPA variantsDistributedSamplerfor litdata streams, which already shard internally, avoiding double-shardinglast.ckpt(weights) and_resume.pt(optimizer/scheduler/epoch/step) atomically, and resumes when both exist. If onlylast.ckptsurvives, it warm-resumes from the weights and infers the epoch frommetrics_history.jsonrather than restarting at zeromodality_0/1/...,total,scale_*) tometrics_history.jsonand W&B, main process onlyval_dataloaderandepoch_callbackfor held-out scoring and pruning, used by the pretraining Optuna sweepScripts and runners
scripts/pretrain_ssl.py:--arch/--method/--task, standardized 128/2/4 defaults,exp_name = {arch}_{method}_{task}_seed{seed}.--task notes_onlyroutes throughNotesLabsMIMIC4(include_labs=False)for a leakage-free notes-only variantscripts/run_full_pretrain.py: replays tuned hyperparameters from abest_paramsJSON, launches env-resolvedtorchrunrather than a bare PATH lookup, and accepts--extrato override tuned values (for example--batch-size 32 --grad-accumulation-steps 2to fit a smaller GPU at the same effective batch size)scripts/run_full_pretrain_local.sh/run_fullpt_condor.sh/scripts/slurm/full_pretrain_cc.sh: skip an encoder only whenmetrics_history.jsonshows a full 50 epochs, so partial runs resume instead of being skipped. SetsPYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueto avoid fragmentation OOMscripts/kill_ddp.sh: top-down kill cascade, needed because torchrun restarts killed workersW&B
pyhealth/_wandb.pyis env-gated onWANDB_PROJECTand swallows all exceptions, so tracking never takes down training. Runs are named{arch}_{method}_{task}_seed{N}and taggedkind:pretrain,stage:pretrain,bb:{arch},mod:{task},method:{method}. Metrics are namespaced intoval/,test/,best/,loss/,sys/so the dashboard groups into a few sections instead of dozens of flat keys.Depends on
The foundation edits (
UnifiedMultimodalEmbeddingModelwithtext_finetune_mode,TransformerLayer,MambaBlock/JambaLayer). Land those first, or fold them in. The Optuna PR sits on top of this one.Testing
tests/test_pretrain.pyandtests/test_pretrain_backbones.py, roughly 680 lines covering masking (random and block, contiguity, floor enforcement), MAE / SimMIM / I-JEPA / V-JEPA forward, backward and loss, RoPE extrapolation, multi-scale sampling, EMA schedule, token vs unified targets, checkpoint loading into the downstream Transformer, and per-modality decoding, plus regressions for degenerate-prediction and empty-batch edge cases. Backbone contract (shape, backward, padding isolation) is validated for all three architectures.Full 50-epoch pretraining has been run end to end on MIMIC-IV across the architecture x method matrix for three modality combinations (notes_labs, notes_only, labs_only), 26 of 27 encoders complete at the time of writing, with healthy per-modality validation losses and no representation collapse.